// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>

using namespace kumo::pollux::functions::aggregate::test;
using namespace kumo::pollux::plan;

namespace kumo::pollux::aggregate::test {
    namespace {
        template<typename T>
        struct ApproxMostFrequentTest : AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
            }

            std::shared_ptr<FlatVector<int> > makeGroupKeys() {
                return make_flat_vector<int>(3, [](auto row) { return row; });
            }

            std::shared_ptr<FlatVector<int> > makeKeys() {
                return make_flat_vector<int>(
                    1000, [](auto row) { return static_cast<int>(std::sqrt(row)) % 3; });
            }

            std::shared_ptr<FlatVector<T> > makeValues() {
                return make_flat_vector<T>(1000, [](auto row) { return std::sqrt(row); });
            }

            std::shared_ptr<FlatVector<T> > makeValuesWithNulls() {
                auto values = makeValues();
                for (int i = 0; i < values->size(); ++i) {
                    if (static_cast<int>(std::sqrt(i)) % 3 == 0) {
                        values->set_null(i, true);
                    }
                }
                return values;
            }

            MapVectorPtr makeGlobalExpected() {
                return make_map_vector<T, int64_t>({{{30, 61}, {29, 59}, {28, 57}}});
            }

            MapVectorPtr makeGroupedExpected() {
                return make_map_vector<T, int64_t>(
                    {
                        {{24, 49}, {27, 55}, {30, 61}},
                        {{22, 45}, {25, 51}, {28, 57}},
                        {{23, 47}, {26, 53}, {29, 59}}
                    });
            }

            MapVectorPtr makeEmptyGroupExpected() {
                auto expected = makeGroupedExpected();
                expected->set_null(0, true);
                return expected;
            }
        };

        template<>
        std::shared_ptr<FlatVector<StringView> >
        ApproxMostFrequentTest<StringView>::makeValues() {
            std::string s[32];
            for (int i = 0; i < 32; ++i) {
                s[i] = std::to_string(i);
            }
            return make_flat_vector<StringView>(1000, [&](auto row) {
                return StringView(s[static_cast<int>(std::sqrt(row))]);
            });
        }

        template<>
        MapVectorPtr ApproxMostFrequentTest<StringView>::makeGlobalExpected() {
            return make_map_vector<StringView, int64_t>(
                {{{"30", 61}, {"29", 59}, {"28", 57}}});
        }

        template<>
        MapVectorPtr ApproxMostFrequentTest<StringView>::makeGroupedExpected() {
            return make_map_vector<StringView, int64_t>(
                {
                    {{"24", 49}, {"27", 55}, {"30", 61}},
                    {{"22", 45}, {"25", 51}, {"28", 57}},
                    {{"23", 47}, {"26", 53}, {"29", 59}}
                });
        }

        using ValueTypes = ::testing::Types<int, StringView>;

        TYPED_TEST_SUITE(ApproxMostFrequentTest, ValueTypes);

        TYPED_TEST(ApproxMostFrequentTest, global) {
            auto values = this->makeValues();
            auto expected = this->makeGlobalExpected();
            this->testAggregations(
                {this->make_row_vector({values})},
                {},
                {"approx_most_frequent(3, c0, 31)"},
                {this->make_row_vector({expected})});
        }

        TYPED_TEST(ApproxMostFrequentTest, grouped) {
            auto values = this->makeValues();
            auto keys = this->makeKeys();
            auto groupKeys = this->makeGroupKeys();
            auto expected = this->makeGroupedExpected();
            this->testAggregations(
                {this->make_row_vector({keys, values})},
                {"c0"},
                {"approx_most_frequent(3, c1, 11)"},
                {this->make_row_vector({groupKeys, expected})});
        }

        TYPED_TEST(ApproxMostFrequentTest, emptyGroup) {
            auto values = this->makeValuesWithNulls();
            auto keys = this->makeKeys();
            auto groupKeys = this->makeGroupKeys();
            auto expected = this->makeEmptyGroupExpected();
            this->testAggregations(
                {this->make_row_vector({keys, values})},
                {"c0"},
                {"approx_most_frequent(3, c1, 11)"},
                {this->make_row_vector({groupKeys, expected})});
        }

        using ApproxMostFrequentTestInt = ApproxMostFrequentTest<int>;

        TEST_F(ApproxMostFrequentTestInt, invalidBuckets) {
            auto rootPool = memory::memoryManager()->addRootPool(
                "test-root", 1 << 21, exec::MemoryReclaimer::create());
            auto leafPool = rootPool->addLeafChild("test-leaf");
            auto run = [&](int64_t buckets) {
                auto rows = make_row_vector({
                    make_constant<int64_t>(buckets, buckets),
                    make_flat_vector<int>(buckets, melon::identity),
                    make_constant<int64_t>(buckets, buckets),
                });
                auto plan = PlanBuilder()
                        .values({rows})
                        .singleAggregation({}, {"approx_most_frequent(c0, c1, c2)"})
                        .planNode();
                return exec::test::AssertQueryBuilder(plan).copyResults(leafPool.get());
            };
            ASSERT_EQ(run(10)->size(), 1);
            try {
                run(1 << 19);
                FAIL() << "Expected an exception";
            } catch (const PolluxException &e) {
                EXPECT_EQ(e.errorCode(), error_code::kMemCapExceeded);
            }
        }

        using ApproxMostFrequentTestStringView = ApproxMostFrequentTest<StringView>;

        TEST_F(ApproxMostFrequentTestStringView, stringLifeCycle) {
            std::string s[32];
            for (int i = 0; i < 32; ++i) {
                s[i] = std::string(StringView::kInlineSize, 'x') + std::to_string(i);
            }
            auto values = make_flat_vector<StringView>(1000, [&](auto row) {
                return StringView(s[static_cast<int>(std::sqrt(row))]);
            });
            auto rows = make_row_vector({values});
            auto expected = make_row_vector({
                make_map_vector<StringView, int64_t>(
                    {
                        {
                            {StringView(s[30]), 122},
                            {StringView(s[29]), 118},
                            {StringView(s[28]), 114}
                        }
                    }),
            });
            testReadFromFiles(
                {rows, rows}, {}, {"approx_most_frequent(3, c0, 31)"}, {expected});
        }

        class ApproxMostFrequentTestBoolean : public AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
            }
        };

        TEST_F(ApproxMostFrequentTestBoolean, basic) {
            auto input = make_row_vector({
                make_flat_vector<int32_t>({0, 1, 0, 1, 0, 1, 0, 1}),
                make_flat_vector<bool>(
                    {true, false, true, true, false, false, false, false}),
                make_constant(true, 8),
                make_constant(false, 8),
                make_all_null_flat_vector<bool>(8),
                make_nullable_flat_vector<bool>(
                    {true, false, std::nullopt, true, false, std::nullopt, false, false}),
            });

            auto expected = make_row_vector({
                make_map_vector<bool, int64_t>({
                    {{true, 3}, {false, 5}},
                }),
            });

            testAggregations(
                {input}, {}, {"approx_most_frequent(3, c1, 31)"}, {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({0, 1}),
                make_map_vector<bool, int64_t>({
                    {{true, 2}, {false, 2}},
                    {{true, 1}, {false, 3}},
                }),
            });

            testAggregations(
                {input}, {"c0"}, {"approx_most_frequent(3, c1, 31)"}, {expected});

            // All 'true'.
            expected = make_row_vector({make_map_vector<bool, int64_t>({{{true, 8}}})});
            testAggregations(
                {input}, {}, {"approx_most_frequent(3, c2, 31)"}, {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({0, 1}),
                make_map_vector<bool, int64_t>({
                    {{true, 4}},
                    {{true, 4}},
                }),
            });

            testAggregations(
                {input}, {"c0"}, {"approx_most_frequent(3, c2, 31)"}, {expected});

            // All 'false'.
            expected = make_row_vector({make_map_vector<bool, int64_t>({{{false, 8}}})});
            testAggregations(
                {input}, {}, {"approx_most_frequent(3, c3, 31)"}, {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({0, 1}),
                make_map_vector<bool, int64_t>({
                    {{false, 4}},
                    {{false, 4}},
                }),
            });

            testAggregations(
                {input}, {"c0"}, {"approx_most_frequent(3, c3, 31)"}, {expected});

            // All nulls.
            expected = make_row_vector({
                BaseVector::create_null_constant(MAP(BOOLEAN(), BIGINT()), 1, pool()),
            });
            testAggregations(
                {input}, {}, {"approx_most_frequent(3, c4, 31)"}, {expected});

            // Some nulls.
            expected = make_row_vector({
                make_map_vector<bool, int64_t>({
                    {{true, 2}, {false, 4}},
                }),
            });

            testAggregations(
                {input}, {}, {"approx_most_frequent(3, c5, 31)"}, {expected});

            expected = make_row_vector({
                make_flat_vector<int32_t>({0, 1}),
                make_map_vector<bool, int64_t>({
                    {{true, 1}, {false, 2}},
                    {{true, 1}, {false, 2}},
                }),
            });

            testAggregations(
                {input}, {"c0"}, {"approx_most_frequent(3, c5, 31)"}, {expected});
        }

        class ApproxMostFrequentTestJson : public AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
            }
        };

        TEST_F(ApproxMostFrequentTestJson, basic) {
            // JSON strings as input
            std::vector<std::string> jsonStrings = {
                "{\"type\": \"store\"}",
                "{\"type\": \"fruit\"}",
                "{\"type\": \"fruit\"}",
                "{\"type\": \"book\"}",
                "{\"type\": \"store\"}",
                "{\"type\": \"fruit\"}"
            };

            auto inputVector = make_flat_vector<StringView>(
                static_cast<vector_size_t>(jsonStrings.size()),
                [&](auto row) { return StringView(jsonStrings[row]); });

            MapVectorPtr expectedMap = make_map_vector<StringView, int64_t>(
                {
                    {
                        {StringView("{\"type\": \"fruit\"}"), 3},
                        {StringView("{\"type\": \"store\"}"), 2}
                    }
                });
            auto expected = make_row_vector({{expectedMap}});

            testAggregations(
                {make_row_vector({inputVector})},
                {},
                {"approx_most_frequent(2, c0, 31)"},
                {expected});
        }
    } // namespace
} // namespace kumo::pollux::aggregate::test
